"""
This implements a morse potential that is ready-to-be-compiled with shedskin.
If you have shedskin installed and working, you can run the commands:

shedskin -e morse.py
make

and it will generate morse.so, which will be imported automatically.
"""


import math



def energy(r, d, a, r0):
    u = 0
    a2 = a * 2
    l = len(r) / 3
    d = 0.0
    diff = 0.0
    for i in range(l - 1):
        for j in range(i + 1, l):
            d = math.sqrt((r[i * 3 + 0] - r[j * 3 + 0])**2 + 
                          (r[i * 3 + 1] - r[j * 3 + 1])**2 + 
                          (r[i * 3 + 2] - r[j * 3 + 2])**2)
            diff = d - r0
            u += d * (math.exp(-a2 * diff) - 2 * math.exp(-a * diff))
    return u



def force(r, d, a, r0):
    a2 = a * 2
    l = len(r) / 3
    f = [0.0, 0.0, 0.0] * l
    for i in range(l - 1):
        for j in range(i + 1, l):
            dx = r[i * 3 + 0] - r[j * 3 + 0]
            dy = r[i * 3 + 1] - r[j * 3 + 1]
            dz = r[i * 3 + 2] - r[j * 3 + 2]
            d = math.sqrt(dx*dx + dy*dy + dz*dz)
            invd = 1.0 / d
            diff = d - r0
            fmag = d * (-a2 * math.exp(-a2*diff) + a2 * math.exp(-a*diff))
            fx = fmag * dx * invd
            fy = fmag * dy * invd
            fz = fmag * dz * invd
            f[i * 3 + 0] -= fx
            f[i * 3 + 1] -= fy
            f[i * 3 + 2] -= fz
            f[j * 3 + 0] += fx
            f[j * 3 + 1] += fy
            f[j * 3 + 2] += fz
    return f
        


if __name__=="__main__":
    r = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
    energy(r, d = 0.2703, a = 1.1646, r0 = 3.253)
    force(r, d = 0.2703, a = 1.1646, r0 = 3.253) 
    
    
    
    
    
    
    
    

    


